﻿from scipy.signal import butter, filtfilt, lfilter_zi, lfilter
import numpy as np
from tqdm import tqdm
from LoadBCICompDataSet import *


class BCIButterworthFilter:
    def __init__(self, order=3, highFrequency=10, lowFrequency=0.5, sample_frequency=240):
        # 滤波器参数
        self.order = order
        self.highFrequency = highFrequency
        self.lowFrequency = lowFrequency
        # 计算归一化截止频率
        '''
        240是BCI competition脑机接口比赛中的采样率
        '''
        self.sample_frequency = sample_frequency
        nyquist_frequency = 0.5 * self.sample_frequency
        normalize_highFrequency = highFrequency / nyquist_frequency
        normalize_lowFrequency = lowFrequency / nyquist_frequency
        # 使用巴特沃斯滤波器设计滤波器系数
        self.b, self.a = butter(order, [normalize_lowFrequency, normalize_highFrequency], btype='bandpass')

    def filtByffilt(self, respones):
        '''
        这里过滤原始数据，就是原始的五维张量
        '''
        # acc_y_filtered = filtfilt(b, a, respones[:,6])
        for x in tqdm(range(np.size(respones, 0))):
            for y in range(np.size(respones, 1)):
                for z in range(np.size(respones, 2)):
                    for w in range(np.size(respones, 4)):
                        respones[x, y, z, :, w] = filtfilt(self.b, self.a, respones[x, y, z, :, w])
        return respones

    def filtByAveffilt(self, respones):
        '''
        这里过滤的是取次数平均后的数据
        '''
        # acc_y_filtered = filtfilt(b, a, respones[:,6])
        for x in tqdm(range(np.size(respones, 0))):
            for y in range(np.size(respones, 2)):
                for z in range(np.size(respones, 3)):
                    respones[x, y, :, z] = filtfilt(self.b, self.a, respones[x, y, :, z])
        return respones

    def filtBylfilt(self, respones):
        '''
        这里过滤原始数据，就是原始的五维张量
        '''
        # 初始化滤波器状态
        zi = lfilter_zi(self.b, self.a)
        # acc_y_filtered_2 = np.zeros_like(respones["colAddRowNum"],respones["ColAndRowAppear"],respones["Windows"],respones["Channels"])#colAddRowNum,15,windows,channels
        # for index in tqdm(range(len(respones))):
        #     raw_value = respones[index,6]
        #     filtered,zi = lfilter(b, a, [raw_value], zi=zi)
        #     acc_y_filtered_2[index] = filtered
        for x in tqdm(range(np.size(respones, 0))):
            for y in range(np.size(respones, 1)):
                for z in range(np.size(respones, 3)):
                    for w in range(np.size(respones, 3)):
                        respones[x, y, z, :, w], zi = lfilter(self.b, self.a, respones[x, y, z, :, w], zi=zi)
        return respones

    def filtByAvelfilt(self, respones):
        '''
        这里过滤的是取次数平均后的数据
        '''
        # 初始化滤波器状态
        zi = lfilter_zi(self.b, self.a)
        # acc_y_filtered_2 = np.zeros_like(respones["colAddRowNum"],respones["ColAndRowAppear"],respones["Windows"],respones["Channels"])#colAddRowNum,15,windows,channels
        # for index in tqdm(range(len(respones))):
        #     raw_value = respones[index,6]
        #     filtered,zi = lfilter(b, a, [raw_value], zi=zi)
        #     acc_y_filtered_2[index] = filtered
        for x in tqdm(range(np.size(respones, 0))):
            for y in range(np.size(respones, 2)):
                for z in range(np.size(respones, 3)):
                    respones[x, y, :, z], zi = lfilter(self.b, self.a, respones[x, y, :, z], zi=zi)
        return respones


if __name__ == "__main__":
    # 输入数据
    data = LoadBCICompDataSet(r"E:\文件\Documents\Python\大创\DataSet\BCI_Comp_III_Wads_2004\Subject_A_Train.mat")
    respones = data.getResponses()
    filter = BCIButterworthFilter()
    ans = filter.filtByffilt(respones)
    print("\n")
    print(ans.shape)
    print(respones[0, 0, 0, :0])
    print(ans[0, 0, :, 0])
